import os
import json
from dotenv import load_dotenv
from langchain.agents import initialize_agent, Tool, AgentType
from langchain_openai import ChatOpenAI  # ✅ 修复1：使用新的导入路径

# 加载环境变量
load_dotenv()

# 检查 API Key
if not os.getenv("OPENAI_API_KEY"):
    raise RuntimeError("请先在 .env 文件中配置 OPENAI_API_KEY")

# ✅ 修复3：直接返回易读字符串，避免 LLM 解析 JSON
def get_weather(city: str) -> str:
    weather_data = {
        "beijing": {
            "location": "北京",
            "rain_probability": 10,
            "description": "晴朗"
        },
        "shenzhen": {
            "location": "深圳",
            "rain_probability": 90,
            "description": "雷阵雨"
        }
    }
    city_key = city.lower()
    if city_key in weather_data:
        w = weather_data[city_key]
        return f"{w['location']}当前{w['description']}，降雨概率{w['rain_probability']}%"
    return f"{city} 的天气信息暂不可用"

# 定义工具
tools = [
    Tool(
        name="GetWeather",
        func=get_weather,
        description=(
            "用于查询指定城市的天气信息，包括降雨概率。"
            "输入参数为城市名称（字符串）。"
        )
    )
]

# 初始化 LLM
llm = ChatOpenAI(model="gpt-3.5-turbo", temperature=0)  # ✅ model 参数名已更新

# 初始化 Agent
agent = initialize_agent(
    tools=tools,
    llm=llm,
    agent=AgentType.CHAT_ZERO_SHOT_REACT_DESCRIPTION,
    verbose=True,
    handle_parsing_errors=True  # ✅ 可选：自动处理解析错误
)

# 运行 Agent
if __name__ == "__main__":
    task = "查找深圳的天气，然后用一句话告诉我出门要不要带伞"
    result = agent.run(task)
    print("\n最终建议:", result)